#include<stdio.h>
#include<cuda.h>
#include<cuda_runtime.h>

#define BLOCK_NUM 4  
#define THREAD_NUM 4
#define R_SIZE BLOCK_NUM * THREAD_NUM
#define M_SIZE R_SIZE * R_SIZE

__global__ void mat_mul(float *mat1, float *mat2, float *result) {
    const int bid = blockIdx.x;
    const int tid = threadIdx.x;
    const int row = bid * THREAD_NUM + tid;
    for (int c = 0; c < R_SIZE; c++) {
        for (int n = 0; n < R_SIZE; n++) {
            result[row*R_SIZE+c] += mat1[row*R_SIZE+n] * mat2[n*R_SIZE+c];
        }
    }
}

int main(int argc, char *argv[]) {
    float *mat1, *mat2, *result;
    float *g_mat1, *g_mat2, *g_mat_result;
    
    // 1-dim NxN vector to represent 2-dim (N, N) matrix
    mat1 = (float*) malloc(M_SIZE * sizeof(float));
    mat2 = (float*) malloc(M_SIZE * sizeof(float));
    result = (float*) malloc(M_SIZE * sizeof(float));
    printf("M_SIZE:%d\n", M_SIZE);
    // init matrices
    for (int i = 0; i < M_SIZE; i++) {
        mat1[i] = (float)(rand() % 10);
        mat2[i] = (float)(rand() % 10);
        result[i] = 0;
    }
    cudaMalloc((void **)&g_mat1, sizeof(float) * M_SIZE);
    cudaMalloc((void **)&g_mat2, sizeof(float) * M_SIZE);
    cudaMalloc((void **)&g_mat_result, sizeof(float) * M_SIZE);
    cudaMemcpy(g_mat1, mat1, sizeof(float) * M_SIZE, cudaMemcpyHostToDevice);
    cudaMemcpy(g_mat2, mat2, sizeof(float) * M_SIZE, cudaMemcpyHostToDevice);
    mat_mul<<<BLOCK_NUM, THREAD_NUM>>>(g_mat1, g_mat2, g_mat_result);
    cudaMemcpy(result, g_mat_result, sizeof(float) * M_SIZE, cudaMemcpyDeviceToHost);
    printf("res[0]:%d\n", result[0]);
}
